import math
import numpy as np
import matplotlib.pyplot as plt
from Create_Environment import Generate_Walls, Generate_Obstacles, Generate_Obstacle_Centers, create_grid_map, Generate_GoalPos, update_vehicle, initialize_plot


def calculate_circles_for_rectangle(x_c, y_c, w, h, n_discs):
    """
        Automatically determine the orientation of a rectangle and calculate multiple circles arranged along its central axis.
    
        Parameters:
            x_c, y_c: Coordinates of the rectangle's center
            w, h: Width and height of the rectangle
            n_discs: Number of circles
    
        Returns:
            circles: Each circle's center and radius, formatted as [(cx1, cy1, r1), ...]
    """
    # Determine the orientation of the rectangle
    is_horizontal = w >= h  
    
    # Calculate the offset of the circle
    if is_horizontal:
        # Horizontal rectangle: Circle centers are evenly distributed along the x direction.
        disc_offset = np.linspace(-w / 2 + w / (2 * n_discs), w / 2 - w / (2 * n_discs), n_discs)
        circles = []
        radius = math.sqrt((w / n_discs / 2) ** 2 + (h / 2) ** 2)  # radius
        for offset in disc_offset:
            cx = x_c + offset  # Circle center x-coordinate
            cy = y_c           # Circle center y-coordinate
            circles.append((cx, cy, radius))
    else:
        # Vertical rectangle: Circle centers are evenly distributed along the y direction.
        disc_offset = np.linspace(-h / 2 + h / (2 * n_discs), h / 2 - h / (2 * n_discs), n_discs)
        circles = []
        radius = math.sqrt((h / n_discs / 2) ** 2 + (w / 2) ** 2)  # 圆的半径
        for offset in disc_offset:
            cx = x_c           # Circle center x-coordinate
            cy = y_c + offset  # Circle center y-coordinate
            circles.append((cx, cy, radius))
    
    return circles, is_horizontal

def generate_circles_from_vehicle(pos0, a, w, h, n_discs):
    """
        Generate the centers and radius of covering circles based on the rear axle position, heading angle, geometric center offset, dimensions, and the number of circles.
    
        Parameters:
            pos0: np.array([x, y, theta]), the vehicle's rear axle center position (x, y) and heading angle theta (in radians).
            a: Distance from the rear axle to the geometric center of the vehicle.
            w: Vehicle width.
            h: Vehicle length.
            n_discs: Number of covering circles.
    
        Returns:
            circles: A list of each circle's center (x, y) and radius r, formatted as [(cx1, cy1, r1), ...].
    """
    
    # Position of the rear axle and orientation
    x_rear, y_rear, theta = pos0  

    # Calculate geometry center of the vehicle 
    x_center = x_rear + a * math.cos(theta)
    y_center = y_rear + a * math.sin(theta)

    # Calculate radius 
    disc_radius = math.sqrt((h / n_discs / 2) ** 2 + (w / 2) ** 2)

    # Circle centers are evenly distributed along the longitudinal axis of the geometric center.
    disc_offsets = np.linspace(
        -h / 2 + h / (2 * n_discs),
        h / 2 - h / (2 * n_discs),
        n_discs
    )

    circles = []
    for offset in disc_offsets:
        # Circle center in the local coordinate system
        cx_local = offset
        cy_local = 0

        # Circle center in the global coordinate system
        cx_global = x_center + cx_local * math.cos(theta) - cy_local * math.sin(theta)
        cy_global = y_center + cx_local * math.sin(theta) + cy_local * math.cos(theta)

        # Add result to the list
        circles.append((cx_global, cy_global, disc_radius))
    
    return circles

def plot_circles(ax, Circles_Dict, Type):
    """
        Draw circles from a list.
        
        Parameters:
            ax: Matplotlib Axes object for plotting.
            Circles_Dict: A dictionary including information about circles.
            Type: Circle type (Obstacle or Vehicle)
    
    """
    
    for circles in Circles_Dict:
        for cx, cy, r in circles:
            if Type == "Obstacle":
                # Obstacle in orange
                circle = plt.Circle((cx, cy), r, color='orange', fill=False, linestyle='--', label='Circle')
                ax.add_patch(circle)
            elif Type == "Vehicle":
                # Vehicle in red
                circle = plt.Circle((cx, cy), r, color='red', fill=False, linestyle='--', label='Circle')
                ax.add_patch(circle)

def Generate_Obstacle_Circles(Wall_Dict, Obs_Dict):
    """
        Generate several circles for static obstacles
        
        Parameters:
            Wall_Dict: A dictionary containing some detailed information about walls
            Obs_Dict: A dictionary containing some detailed information about static obstacles
        
        Return:
            Circles_Dict: A dictionary containing some detailed information about circles
    
    """
    
    # Create circles for walls
    Circles_Dict = []
    for Wall in  Wall_Dict:
        pos = Wall['geometry']['position']
        h = Wall['geometry']['width']
        w = Wall['geometry']['length']
        x_c = pos[0]
        y_c = pos[1]

        if max(h,w) == 80:
            n_discs = 80
        elif max(h,w) == 50:
            n_discs = 50
        elif max(h,w) == 25:
            n_discs = 25
        else:
            n_discs = 10
        circles, _ = calculate_circles_for_rectangle(x_c, y_c, w, h, n_discs)
        # Circles_Dict.append(circles)
        
    # Create circles for static obstacles  
    for Obstacle in  Obs_Dict:
        pos = Obstacle['geometry']['position']
        h = Obstacle['geometry']['width']
        w = Obstacle['geometry']['length']
        x_c = pos[0]
        y_c = pos[1]

        n_discs = 2
        circles, _ = calculate_circles_for_rectangle(x_c, y_c, w, h, n_discs)
        Circles_Dict.append(circles)
            
    return Circles_Dict

# Testing Code
if __name__ == "__main__":
    # Create environment
    length = 50
    width = 80
    Walls, Wall_Dict = Generate_Walls(length, width, Type='Parking Lot')
    
    # Create obstacles
    Obs_centers = Generate_Obstacle_Centers(length, width, n_obstacles=12, Type='Parking Lot')
    Obstacles , Obs_Dict = Generate_Obstacles(Obs_centers, Type='Parking Lot')
    
    # Generate circles
    Circles_Dict = Generate_Obstacle_Circles(Wall_Dict, Obs_Dict)
    
    # Define Grid Resolution
    grid_size = 0.1
    grid_map, min_x, max_x, min_y, max_y = create_grid_map(Wall_Dict, Obs_Dict, grid_size=grid_size)
    
    # Generate a vehicle
    pos0 = [-20, -35, np.pi/6]
    a = 1.41
    w = 2.1
    h = 4.76
    Vehicle_Circle = []
    Vehicle_Circle.append(generate_circles_from_vehicle(pos0, a, w, h, n_discs=2))
    
    # Visualize the result
    fig, ax = initialize_plot(grid_map, min_x, max_x, min_y, max_y)
    plot_circles(ax, Circles_Dict,Type='Obstacle')
    update_vehicle(ax, pos0)
    plot_circles(ax, Vehicle_Circle,Type='Vehicle')
